import tushare as ts
import pandas as pd
import os
from tqdm import tqdm

# 配置Token
ts.set_token('gx03013e909f633ecb66722df66b360f070426613316ebf06ecd3482')
pro = ts.pro_api()

# 路径配置
STOCK_LIST_PATH = r"F:\量化投资程序\机器算法准备\综合修改版\酒类行业明细.xlsx"
OUTPUT_DIR = r"F:\上市公司综合数据"
OUTPUT_FILE = os.path.join(OUTPUT_DIR, "酒类行业上市公司综合数据.csv")
REPORT_YEAR = 2023  # 统计年份


def load_stock_codes():
    """加载股票代码列表"""
    try:
        df = pd.read_excel(STOCK_LIST_PATH)
        print(f"成功加载{len(df)}支酒类行业股票")
        return df['ts_code'].unique().tolist()
    except Exception as e:
        print(f"股票列表加载失败: {str(e)}")
        return []


def get_exchange(code):
    """获取交易所代码"""
    return "SZSE" if code.endswith(".SZ") else "SSE"


def fetch_company_info(code):
    """获取公司基本信息"""
    try:
        return pro.stock_company(
            exchange=get_exchange(code),
            ts_code=code,
            fields="ts_code,com_name,employees,province,main_business"
        )
    except Exception as e:
        print(f"{code} 公司信息获取失败: {str(e)[:50]}")
        return pd.DataFrame()


def fetch_financial_data(code):
    """获取财务数据"""
    try:
        # 获取最新年报数据
        report_date = f"{REPORT_YEAR}1231"
        df = pro.balancesheet_vip(
            ts_code=code,
            period=report_date,
            fields="ts_code,end_date,total_assets,total_hldr_eqy_exc_min_int,inventories,accounts_receiv"
        )

        if not df.empty:
            # 计算财务指标
            df["资产负债率"] = (1 - df["total_hldr_eqy_exc_min_int"] / df["total_assets"]).round(4)
            df["存货(亿元)"] = (df["inventories"] / 1e8).round(2)
            df["应收账款(亿元)"] = (df["accounts_receiv"] / 1e8).round(2)
        return df
    except Exception as e:
        print(f"{code} 财务数据获取失败: {str(e)[:50]}")
        return pd.DataFrame()


def process_stock(code):
    """处理单支股票"""
    # 获取数据
    company_df = fetch_company_info(code)
    finance_df = fetch_financial_data(code)

    # 合并数据
    if not company_df.empty and not finance_df.empty:
        merged_df = pd.merge(company_df, finance_df, on="ts_code")
        return merged_df[[
            "ts_code", "com_name", "province", "main_business",
            "employees", "资产负债率", "存货(亿元)", "应收账款(亿元)",
            "end_date"
        ]].rename(columns={
            "ts_code": "股票代码",
            "com_name": "公司名称",
            "province": "所在省份",
            "main_business": "主营业务",
            "employees": "员工人数",
            "end_date": "报告日期"
        })
    return pd.DataFrame()


def batch_process():
    """批量处理主程序"""
    # 创建输出目录
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # 加载股票列表
    stock_list = load_stock_codes()
    if not stock_list:
        return

    # 处理数据
    all_data = []
    error_log = []

    # 添加进度条
    pbar = tqdm(stock_list, desc="处理进度", unit="股票")

    for code in pbar:
        pbar.set_postfix_str(code)
        try:
            df = process_stock(code)
            if not df.empty:
                all_data.append(df)
            else:
                error_log.append({"股票代码": code, "错误原因": "数据获取失败"})
        except Exception as e:
            error_log.append({"股票代码": code, "错误原因": str(e)[:50]})

    # 保存结果
    if all_data:
        final_df = pd.concat(all_data, ignore_index=True)

        # 添加统计年份列
        final_df["统计年份"] = REPORT_YEAR

        # 列排序优化
        column_order = [
            "股票代码", "公司名称", "所在省份", "统计年份",
            "员工人数", "资产负债率", "存货(亿元)", "应收账款(亿元)",
            "主营业务", "报告日期"
        ]
        final_df = final_df[column_order]

        final_df.to_csv(OUTPUT_FILE, index=False, encoding="utf_8_sig")
        print(f"\n成功处理 {len(all_data)}/{len(stock_list)} 支股票")
        print(f"结果文件: {OUTPUT_FILE}")
    else:
        print("\n未获取到有效数据")

    # 保存错误日志
    if error_log:
        error_df = pd.DataFrame(error_log)
        error_file = os.path.join(OUTPUT_DIR, "数据采集异常记录.csv")
        error_df.to_csv(error_file, index=False, encoding="utf_8_sig")
        print(f"异常记录已保存至: {error_file}")


if __name__ == "__main__":
    batch_process()
    print("\n处理完成，请检查输出目录")

